#!/usr/bin/env python
# -*- coding=utf-8 -*-
import sys
import os
import re
import subprocess
from textwrap import dedent

def parse_sample_list(sample_list):
    """
    parse sample list and get the lib path easily
    """
    list_dict = {}
    with open(sample_list,'r') as indata:
        title = indata.readline()
        for line in indata:
            line_list = line.strip().split('\t')
            lane = line_list[0]
            sample_name = line_list[1]
            libid = line_list[3]
            path = line_list[6]
            list_dict[sample_name] = {
                "stat1":libid + "_L{lane}_1.adap.stat".format(lane=lane),
                "stat2":libid + "_L{lane}_2.adap.stat".format(lane=lane),
                "fq1": libid + "_L{lane}_1.fq.gz".format(lane=lane),
                "fq2": libid + "_L{lane}_2.fq.gz".format(lane=lane),
                "lib_path" : os.path.join(path,libid),
                "libid":libid
                }
    return list_dict

def truncate_fq(list_dict,odir,script_dir):
    """
    1) get the reads number 
    2) if data is larger than 12G truncating it to 12G 
    3) if data is less than 12G doing nothing
    """
    for sample in list_dict:
        stat1 = list_dict[sample]["stat1"]
        stat2 = list_dict[sample]["stat2"]
        libid = list_dict[sample]["libid"]
        lib_path = list_dict[sample]["lib_path"]
        fq1 = list_dict[sample]["fq1"]
        fq2 = list_dict[sample]["fq2"]
        stat1_path = os.path.join(lib_path,stat1)
        fq1_path = os.path.join(lib_path,fq1)
        fq2_path = os.path.join(lib_path,fq2)
        with open(stat1_path,'r') as indata:
            for line in indata:
                resp = re.match(r"^total_reads: (\d+)$",line)
                if resp is not None:
                    reads_number = int(resp.group(1))
                    data_size = float(int(reads_number)*150*2)/1000000000
                    rest_reads = float((data_size - 12)*1000000000)/(150*2)
                    ofile1 = os.path.join(odir,libid,fq1)
                    ofile2 = os.path.join(odir,libid,fq2)
                    if not os.path.exists(os.path.join(odir,libid)):
                        os.makedirs(os.path.join(odir,libid))

                    if rest_reads > 0:
                        print('{sample} data size is {data_size},which is larger than 12G ,truncate it!'.format(**locals()))
                        print(rest_reads)
                        new_reads_number = reads_number - int(rest_reads)
                        print(new_reads_number)
                        cmd = """
                        set -eo pipefail
                        echo start `date +"%F %T"`
                        seqkit head  \\
                            -n {new_reads_number} \\
                            -j 8 \\
                            -o {ofile1} \\
                            {fq1_path}
                            
                        seqkit head  \\
                            -n {new_reads_number} \\
                            -j 8 \\
                            -o {ofile2} \\
                            {fq2_path}

                        echo end `date +"%F %T"`
                        """.format(**locals())
                        oshell = os.path.join(script_dir,"work.{sample}.truncate.sh".format(sample=sample))
                        with open(oshell,'w') as odata:
                            odata.write(dedent(cmd))
                        print('{sample} truncating script write to {oshell}'.format(**locals()))
                    else:
                        print('{sample} data size is {data_size},which is less than 12G ,do nothing!'.format(**locals()))

def main():
    sample_list = sys.argv[1]
    odir = sys.argv[2]
    script_dir = sys.argv[3]
    list_dict = parse_sample_list(sample_list)
    truncate_fq(list_dict,odir,script_dir)

if __name__ == "__main__":
    main()

